
-- set search_path to sm_sc;
-- drop function if exists sm_sc.fv_d_redistr_softmax_dloss_dindepdt(float[], float[], float[]);
create or replace function sm_sc.fv_d_redistr_softmax_dloss_dindepdt
(
  i_depdt                  float[]                 ,                     -- softmax 输出
  i_dloss_ddepdt           float[]                 ,                     -- 此入参传入 dloss/dindepdt, 用于 softmax 直接求取 dloss/ddepdt
  i_indepdt                float[]      default null                     -- softmax 算子的输入，来自上一层算子的输出
)
returns float[]     -- 输出列序与 i_indepdt 枚举值 one_hot 序一致
as
$$
declare 
  v_len               int;
  v_stddev_samp       float;
  v_ret               float[]   := array_fill(null :: float, (select array_agg(array_length(i_dloss_ddepdt, a_no) order by a_no) from generate_series(1, array_ndims(i_dloss_ddepdt)) tb_a(a_no)));
begin
  -- 审计维度
  if current_setting('pg4ml._v_is_debug_check', true) = '1'
  then
    if array_ndims(i_depdt) > 4 or array_ndims(i_indepdt) > 4 or array_ndims(i_dloss_ddepdt) > 4
      or array_dims(i_dloss_ddepdt) <> array_dims(i_depdt) or array_dims(i_dloss_ddepdt) <> array_dims(i_indepdt)
      -- -- or array_length(i_dloss_ddepdt, 1) <> array_length(i_depdt, 1)
      -- -- or array_length(i_dloss_ddepdt, 2) <> array_length(i_depdt, 2)
      -- -- or array_length(i_dloss_ddepdt, 3) <> array_length(i_depdt, 3)
      -- -- or array_length(i_dloss_ddepdt, 4) <> array_length(i_depdt, 4)
      -- -- or array_length(i_dloss_ddepdt, 1) <> array_length(i_indepdt, 1)
      -- -- or array_length(i_dloss_ddepdt, 2) <> array_length(i_indepdt, 2)
      -- -- or array_length(i_dloss_ddepdt, 3) <> array_length(i_indepdt, 3)
      -- -- or array_length(i_dloss_ddepdt, 4) <> array_length(i_indepdt, 4)
    then 
      raise exception 'unmatch ndims or length between i_depdt, i_indepdt and i_dloss_ddepdt.';
    end if;
  end if;

  if i_depdt is null
  then
    i_depdt := sm_sc.fv_redistr_softmax(i_indepdt);
  end if;

  v_ret[ : ] :=
    sm_sc.__fv_d_redistr_softmax_1d_dloss_dindepdt
    (
      (select array[array_agg(a_col)] from unnest(i_depdt) tb_a(a_col)), 
      (select array[array_agg(a_col)] from unnest(i_dloss_ddepdt) tb_a(a_col)), 
      (select array[array_agg(a_col)] from unnest(i_indepdt) tb_a(a_col))
    )
  ; -- ~=` 8
  
  return v_ret;
end
$$
language plpgsql stable
parallel safe
cost 100;
-- -- set search_path to sm_sc;
-- select sm_sc.fv_d_redistr_softmax_dloss_dindepdt
--   (
--     sm_sc.fv_redistr_softmax(array[array[1, 5], array[2, 4], array[3, 3], array[4, 2], array[5, 1]]),
--     (-` array[array[0.0 :: float, 1.0], array[0.0 :: float, 0.0], array[0.0 :: float, 0.0], array[0.0 :: float, 0.0], array[1.0 :: float, 0.0]] /` sm_sc.fv_redistr_softmax(array[array[1, 5], array[2, 4], array[3, 3], array[4, 2], array[5, 1]])),
--     array[array[1, 5], array[2, 4], array[3, 3], array[4, 2], array[5, 1]]
--   );

-- select sm_sc.fv_d_redistr_softmax_dloss_dindepdt
--   (
--     sm_sc.fv_redistr_softmax(array[array[1, 5], array[2, 4], array[3, 3], array[4, 2], array[5, 1]]),
--     (-` array[array[0.0 :: float, 1.0], array[0.0 :: float, 0.0], array[0.0 :: float, 0.0], array[0.0 :: float, 0.0], array[1.0 :: float, 0.0]] /` sm_sc.fv_redistr_softmax(array[array[1, 5], array[2, 4], array[3, 3], array[4, 2], array[5, 1]])),
--     array[array[1, 5], array[2, 4], array[3, 3], array[4, 2], array[5, 1]]
--   );

-- with 
-- cte_arr as 
-- (
--   select 
--     sm_sc.fv_new_rand(array[2, 3, 4]) as a_arr
-- )
-- select 
--   sm_sc.fv_d_redistr_softmax_dloss_dindepdt
--   (
--     sm_sc.fv_redistr_softmax(a_arr)
--   , array[[[0,0,0,0],[0,0,0,0],[0,0,0,0]],[[1,0,0,0],[0,0,0,0],[0,0,0,0]]]
--   , a_arr
--   ) 
-- from cte_arr

-- with 
-- cte_arr as 
-- (
--   select 
--     sm_sc.fv_new_rand(array[2, 3, 4, 2]) as a_arr
-- )
-- select 
--   sm_sc.fv_d_redistr_softmax_dloss_dindepdt
--   (
--     sm_sc.fv_redistr_softmax(a_arr)
--   , array[[[[0,0],[0,0],[0,0],[0,0]],[[0,0],[0,0],[0,0],[0,0]],[[0,0],[0,0],[0,0],[0,0]]],[[[1,0],[0,0],[0,0],[0,0]],[[0,0],[0,0],[0,0],[0,0]],[[0,0],[0,0],[0,0],[0,0]]]]
--   , a_arr
--   ) 
-- from cte_arr


-- ------------------------------------------------------------------------------------------------------
-- drop function if exists sm_sc.fv_d_redistr_softmax_dloss_dindepdt(float[], float[], float[], int[]);
create or replace function sm_sc.fv_d_redistr_softmax_dloss_dindepdt
(
  i_depdt                  float[]                  ,                     -- softmax 输出
  i_dloss_ddepdt           float[]                  ,                     -- 此入参传入 dloss/dindepdt, 用于 softmax 直接求取 dloss/ddepdt
  i_indepdt                  float[]                  ,                     -- softmax 算子的输入，来自上一层算子的输出
  i_cnt_per_grp        int[]
)
returns float[]
as
$$
declare 
  v_cur_y  int               ;
  v_cur_x  int               ;
  v_cur_x3 int               ;
  v_cur_x4 int               ;
  v_ret               float[]   := array_fill(null :: float, (select array_agg(array_length(i_dloss_ddepdt, a_no) order by a_no) from generate_series(1, array_ndims(i_dloss_ddepdt)) tb_a(a_no)));
  
begin
  if current_setting('pg4ml._v_is_debug_check', true) = '1'
  then
    if array_ndims(i_cnt_per_grp) > 1 
    then 
      raise exception 'unsupport ndims of i_cnt_per_grp > 1.';
    elsif array_ndims(i_depdt) > 4 or array_ndims(i_indepdt) > 4 or array_ndims(i_dloss_ddepdt) > 4
      or array_dims(i_dloss_ddepdt) <> array_dims(i_depdt) or array_dims(i_dloss_ddepdt) <> array_dims(i_indepdt)
      -- -- or array_length(i_dloss_ddepdt, 1) <> array_length(i_depdt, 1)
      -- -- or array_length(i_dloss_ddepdt, 2) <> array_length(i_depdt, 2)
      -- -- or array_length(i_dloss_ddepdt, 3) <> array_length(i_depdt, 3)
      -- -- or array_length(i_dloss_ddepdt, 4) <> array_length(i_depdt, 4)
      -- -- or array_length(i_dloss_ddepdt, 1) <> array_length(i_indepdt, 1)
      -- -- or array_length(i_dloss_ddepdt, 2) <> array_length(i_indepdt, 2)
      -- -- or array_length(i_dloss_ddepdt, 3) <> array_length(i_indepdt, 3)
      -- -- or array_length(i_dloss_ddepdt, 4) <> array_length(i_indepdt, 4)
    then 
      raise exception 'unmatch ndims or length between i_depdt, i_indepdt and i_dloss_ddepdt.';
    elsif 
      0 <> any 
      (
        (
          select 
            array_agg(array_length(i_dloss_ddepdt, a_cur_dim) order by a_cur_dim) 
          from generate_series(1, array_ndims(i_dloss_ddepdt)) tb_a_cur_dim(a_cur_dim)
        )
        %` i_cnt_per_grp
      )
    then 
      raise exception 'unperfect i_indepdt''s length for i_cnt_per_grp at some dims';
    end if;
  end if;
  
  if i_dloss_ddepdt is null
  then 
    return null;
    
  elsif i_cnt_per_grp is null 
  then 
    return sm_sc.fv_d_redistr_softmax_dloss_dindepdt(i_depdt, i_dloss_ddepdt, i_indepdt);
    
  elsif array_length(i_cnt_per_grp, 1) = 1
  then 
    for v_cur_y in 1 .. array_length(i_dloss_ddepdt, 1) / i_cnt_per_grp[1]
    loop 
      v_ret
        [(v_cur_y - 1) * i_cnt_per_grp[1] + 1 : v_cur_y * i_cnt_per_grp[1]] 
      := 
        sm_sc.fv_d_redistr_softmax_dloss_dindepdt
        (
          i_depdt
            [(v_cur_y - 1) * i_cnt_per_grp[1] + 1 : v_cur_y * i_cnt_per_grp[1]]
        , i_dloss_ddepdt
            [(v_cur_y - 1) * i_cnt_per_grp[1] + 1 : v_cur_y * i_cnt_per_grp[1]]
        , i_indepdt
            [(v_cur_y - 1) * i_cnt_per_grp[1] + 1 : v_cur_y * i_cnt_per_grp[1]]
        );
    end loop;
    return v_ret;
    
  elsif array_length(i_cnt_per_grp, 1) = 2
  then    
    for v_cur_y in 1 .. array_length(i_dloss_ddepdt, 1) / i_cnt_per_grp[1]
    loop 
      for v_cur_x in 1 .. array_length(i_dloss_ddepdt, 2) / i_cnt_per_grp[2]
      loop 
        v_ret
          [(v_cur_y - 1) * i_cnt_per_grp[1] + 1 : v_cur_y * i_cnt_per_grp[1]]
          [(v_cur_x - 1) * i_cnt_per_grp[2] + 1 : v_cur_x * i_cnt_per_grp[2]] 
        := 
          sm_sc.fv_d_redistr_softmax_dloss_dindepdt
          (
            i_depdt
              [(v_cur_y - 1) * i_cnt_per_grp[1] + 1 : v_cur_y * i_cnt_per_grp[1]]
              [(v_cur_x - 1) * i_cnt_per_grp[2] + 1 : v_cur_x * i_cnt_per_grp[2]]
          , i_dloss_ddepdt
              [(v_cur_y - 1) * i_cnt_per_grp[1] + 1 : v_cur_y * i_cnt_per_grp[1]]
              [(v_cur_x - 1) * i_cnt_per_grp[2] + 1 : v_cur_x * i_cnt_per_grp[2]]
          , i_indepdt
              [(v_cur_y - 1) * i_cnt_per_grp[1] + 1 : v_cur_y * i_cnt_per_grp[1]]
              [(v_cur_x - 1) * i_cnt_per_grp[2] + 1 : v_cur_x * i_cnt_per_grp[2]]
          );
      end loop;
    end loop;
    return v_ret;
    
  elsif array_length(i_cnt_per_grp, 1) = 3
  then
    for v_cur_y in 1 .. array_length(i_dloss_ddepdt, 1) / i_cnt_per_grp[1]
    loop 
      for v_cur_x in 1 .. array_length(i_dloss_ddepdt, 2) / i_cnt_per_grp[2]
      loop 
        for v_cur_x3 in 1 .. array_length(i_dloss_ddepdt, 3) / i_cnt_per_grp[3]
        loop 
          v_ret
            [(v_cur_y - 1) * i_cnt_per_grp[1] + 1 : v_cur_y * i_cnt_per_grp[1]]
            [(v_cur_x - 1) * i_cnt_per_grp[2] + 1 : v_cur_x * i_cnt_per_grp[2]]
            [(v_cur_x3 - 1) * i_cnt_per_grp[3] + 1 : v_cur_x3 * i_cnt_per_grp[3]] 
          := 
            sm_sc.fv_d_redistr_softmax_dloss_dindepdt
            (
              i_depdt
                [(v_cur_y - 1) * i_cnt_per_grp[1] + 1 : v_cur_y * i_cnt_per_grp[1]]
                [(v_cur_x - 1) * i_cnt_per_grp[2] + 1 : v_cur_x * i_cnt_per_grp[2]]
                [(v_cur_x3 - 1) * i_cnt_per_grp[3] + 1 : v_cur_x3 * i_cnt_per_grp[3]]
            , i_dloss_ddepdt
                [(v_cur_y - 1) * i_cnt_per_grp[1] + 1 : v_cur_y * i_cnt_per_grp[1]]
                [(v_cur_x - 1) * i_cnt_per_grp[2] + 1 : v_cur_x * i_cnt_per_grp[2]]
                [(v_cur_x3 - 1) * i_cnt_per_grp[3] + 1 : v_cur_x3 * i_cnt_per_grp[3]]
            , i_indepdt
                [(v_cur_y - 1) * i_cnt_per_grp[1] + 1 : v_cur_y * i_cnt_per_grp[1]]
                [(v_cur_x - 1) * i_cnt_per_grp[2] + 1 : v_cur_x * i_cnt_per_grp[2]]
                [(v_cur_x3 - 1) * i_cnt_per_grp[3] + 1 : v_cur_x3 * i_cnt_per_grp[3]]
            );
        end loop;
      end loop;
    end loop;
    return v_ret;
    
  elsif array_length(i_cnt_per_grp, 1) = 4
  then
    for v_cur_y in 1 .. array_length(i_dloss_ddepdt, 1) / i_cnt_per_grp[1]
    loop 
      for v_cur_x in 1 .. array_length(i_dloss_ddepdt, 2) / i_cnt_per_grp[2]
      loop 
        for v_cur_x3 in 1 .. array_length(i_dloss_ddepdt, 3) / i_cnt_per_grp[3]
        loop 
          for v_cur_x4 in 1 .. array_length(i_dloss_ddepdt, 4) / i_cnt_per_grp[4]
          loop 
            v_ret
              [(v_cur_y - 1) * i_cnt_per_grp[1] + 1 : v_cur_y * i_cnt_per_grp[1]]
              [(v_cur_x - 1) * i_cnt_per_grp[2] + 1 : v_cur_x * i_cnt_per_grp[2]]
              [(v_cur_x3 - 1) * i_cnt_per_grp[3] + 1 : v_cur_x3 * i_cnt_per_grp[3]]
              [(v_cur_x4 - 1) * i_cnt_per_grp[4] + 1 : v_cur_x4 * i_cnt_per_grp[4]] 
            := 
              sm_sc.fv_d_redistr_softmax_dloss_dindepdt
              (
                i_depdt
                  [(v_cur_y - 1) * i_cnt_per_grp[1] + 1 : v_cur_y * i_cnt_per_grp[1]]
                  [(v_cur_x - 1) * i_cnt_per_grp[2] + 1 : v_cur_x * i_cnt_per_grp[2]]
                  [(v_cur_x3 - 1) * i_cnt_per_grp[3] + 1 : v_cur_x3 * i_cnt_per_grp[3]]
                  [(v_cur_x4 - 1) * i_cnt_per_grp[4] + 1 : v_cur_x4 * i_cnt_per_grp[4]]
              , i_dloss_ddepdt
                  [(v_cur_y - 1) * i_cnt_per_grp[1] + 1 : v_cur_y * i_cnt_per_grp[1]]
                  [(v_cur_x - 1) * i_cnt_per_grp[2] + 1 : v_cur_x * i_cnt_per_grp[2]]
                  [(v_cur_x3 - 1) * i_cnt_per_grp[3] + 1 : v_cur_x3 * i_cnt_per_grp[3]]
                  [(v_cur_x4 - 1) * i_cnt_per_grp[4] + 1 : v_cur_x4 * i_cnt_per_grp[4]]
              , i_indepdt
                  [(v_cur_y - 1) * i_cnt_per_grp[1] + 1 : v_cur_y * i_cnt_per_grp[1]]
                  [(v_cur_x - 1) * i_cnt_per_grp[2] + 1 : v_cur_x * i_cnt_per_grp[2]]
                  [(v_cur_x3 - 1) * i_cnt_per_grp[3] + 1 : v_cur_x3 * i_cnt_per_grp[3]]
                  [(v_cur_x4 - 1) * i_cnt_per_grp[4] + 1 : v_cur_x4 * i_cnt_per_grp[4]]
              );
          end loop;
        end loop;
      end loop;
    end loop;
    return v_ret;
  end if;
end
$$
language plpgsql stable
parallel safe
cost 100;

-- do
-- $$
-- declare 
--   v_dloss_ddepdt   float[]  :=   sm_sc.fv_new(0.0, array[4, 6, 9]);
-- begin
--   v_dloss_ddepdt[1][1][1]   :=   1.0;
--   raise notice '%',
--   (
--     with 
--     cte_arr as 
--     (
--       select 
--         sm_sc.fv_new_rand(array[4, 6, 9]) as a_arr
--     )
--     select 
--       sm_sc.fv_d_redistr_softmax_dloss_dindepdt
--       (
--         sm_sc.fv_redistr_softmax(a_arr)
--       , v_dloss_ddepdt
--       , a_arr
--       , array[2, 3, 3]
--       ) 
--     from cte_arr
--   );
-- end
-- $$
-- language plpgsql